package org.example.im.websocket;

import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;
import org.example.im.dto.MessageDTO;
import org.example.im.dto.UserDTO;
import org.example.im.dto.WebSocketMessage;
import org.example.im.entity.User;
import org.example.im.service.MessageService;
import org.example.im.service.UserService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Component
@ServerEndpoint("/ws/chat/{userId}")
public class WebSocketServer {
    
    // 用静态变量保存在线用户的WebSocket连接
    private static final Map<Long, Session> ONLINE_SESSIONS = new ConcurrentHashMap<>();
    
    // 因为WebSocket是多例的，所以使用静态变量引用Spring管理的Bean
    private static UserService userService;
    private static MessageService messageService;
    
    // 注入服务
    @Autowired
    public void setUserService(UserService userService) {
        WebSocketServer.userService = userService;
    }
    
    @Autowired
    public void setMessageService(MessageService messageService) {
        WebSocketServer.messageService = messageService;
    }
    
    // 连接建立时调用
    @OnOpen
    public void onOpen(Session session, @PathParam("userId") Long userId) {
        log.info("用户{}连接到WebSocket服务器", userId);
        
        // 将用户加入在线列表
        ONLINE_SESSIONS.put(userId, session);
        
        try {
            // 获取用户信息
            User user = userService.getUserById(userId);
            UserDTO userDTO = new UserDTO();
            userDTO.setId(user.getId());
            userDTO.setUsername(user.getUsername());
            userDTO.setOnline(true);
            
            // 广播用户上线消息
            WebSocketMessage message = WebSocketMessage.createUserOnlineMessage(userDTO);
            broadcastAllUsers(message);
        } catch (Exception e) {
            log.error("处理用户连接时发生错误", e);
        }
    }
    
    // 连接关闭时调用
    @OnClose
    public void onClose(@PathParam("userId") Long userId) {
        log.info("用户{}断开WebSocket连接", userId);
        
        // 将用户从在线列表移除
        ONLINE_SESSIONS.remove(userId);
        
        try {
            // 获取用户信息
            User user = userService.getUserById(userId);
            UserDTO userDTO = new UserDTO();
            userDTO.setId(user.getId());
            userDTO.setUsername(user.getUsername());
            userDTO.setOnline(false);
            
            // 广播用户下线消息
            WebSocketMessage message = WebSocketMessage.createUserOfflineMessage(userDTO);
            broadcastAllUsers(message);
        } catch (Exception e) {
            log.error("处理用户断开连接时发生错误", e);
        }
    }
    
    // 收到消息时调用
    @OnMessage
    public void onMessage(String message, @PathParam("userId") Long userId) {
        log.info("收到用户{}的消息: {}", userId, message);
        
        try {
            // 解析消息
            WebSocketMessage webSocketMessage = JSON.parseObject(message, WebSocketMessage.class);
            
            if (webSocketMessage.getType() == 0) { // 聊天消息
                MessageDTO messageDTO = JSON.parseObject(JSON.toJSONString(webSocketMessage.getData()), MessageDTO.class);
                messageDTO.setSenderId(userId);
                
                // 保存消息到数据库
                messageService.saveMessage(messageDTO);
                
                // 设置发送者名称
                User sender = userService.getUserById(userId);
                messageDTO.setSenderName(sender.getUsername());
                
                // 创建聊天消息并发送
                WebSocketMessage chatMessage = WebSocketMessage.createChatMessage(messageDTO);
                
                if (messageDTO.getChatType() == 0) { // 私聊
                    // 发送给接收者
                    sendMessage(messageDTO.getReceiverId(), chatMessage);
                    
                    // 发送给发送者（回显）
                    sendMessage(userId, chatMessage);
                } else { // 群聊
                    broadcastAllUsers(chatMessage);
                }
            } else if (webSocketMessage.getType() == 3) { // 心跳检测
                sendMessage(userId, WebSocketMessage.createHeartbeatMessage());
            }
        } catch (Exception e) {
            log.error("处理消息时发生错误", e);
        }
    }
    
    // 发生错误时调用
    @OnError
    public void onError(Session session, Throwable error) {
        log.error("WebSocket发生错误", error);
    }
    
    // 发送消息给指定用户
    public static void sendMessage(Long userId, WebSocketMessage message) {
        Session session = ONLINE_SESSIONS.get(userId);
        if (session != null && session.isOpen()) {
            try {
                session.getBasicRemote().sendText(JSON.toJSONString(message));
            } catch (IOException e) {
                log.error("发送消息给用户{}失败", userId, e);
            }
        }
    }
    
    // 广播消息给所有在线用户
    public static void broadcastAllUsers(WebSocketMessage message) {
        ONLINE_SESSIONS.forEach((userId, session) -> {
            if (session.isOpen()) {
                try {
                    session.getBasicRemote().sendText(JSON.toJSONString(message));
                } catch (IOException e) {
                    log.error("广播消息给用户{}失败", userId, e);
                }
            }
        });
    }
    
    // 检查用户是否在线
    public static boolean isUserOnline(Long userId) {
        return ONLINE_SESSIONS.containsKey(userId);
    }
    
    // 初始化
    @PostConstruct
    public void init() {
        log.info("WebSocketServer已初始化");
    }
} 